//----------------------------------------------------------------------
//	File:		kd_pr_search.cc
//	Programmer:	Sunil Arya and David Mount
//	Last modified:	03/04/98 (Release 0.1)
//	Description:	Priority search for kd-trees
//----------------------------------------------------------------------
// Copyright (c) 1997-1998 University of Maryland and Sunil Arya and David
// Mount.  All Rights Reserved.
// 
// This software and related documentation is part of the 
// Approximate Nearest Neighbor Library (ANN).
// 
// Permission to use, copy, and distribute this software and its 
// documentation is hereby granted free of charge, provided that 
// (1) it is not a component of a commercial product, and 
// (2) this notice appears in all copies of the software and
//     related documentation. 
// 
// The University of Maryland (U.M.) and the authors make no representations
// about the suitability or fitness of this software for any purpose.  It is
// provided "as is" without express or implied warranty.
//----------------------------------------------------------------------

#include "ANN.h"
#include "ANNx.h"			// all ANN includes
#include "ANNperf.h"		// performance evaluation
#include "kd_pr_search.h"		// kd priority search declarations

//----------------------------------------------------------------------
//  Approximate nearest neighbor searching by priority search.
//	The kd-tree is searched for an approximate nearest neighbor.
//	The point is returned through one of the arguments, and the
//	distance returned is the SQUARED distance to this point.
//
//	The method used for searching the kd-tree is called priority
//	search.  (It is described in Arya and Mount, ``Algorithms for
//	fast vector quantization,'' Proc. of DCC '93: Data Compression
//	Conference}, eds. J. A.  Storer and M. Cohn, IEEE Press, 1993,
//	381--390.)
//
//	The cell of the kd-tree containing the query point is located,
//	and cells are visited in increasing order of distance from the
//	query point.  This is done by placing each subtree which has
//	NOT been visited in a priority queue, according to the closest
//	distance of the corresponding enclosing rectangle from the
//	query point.  The search stops when the distance to the nearest
//	remaining rectangle exceeds the distance to the nearest point
//	seen by a factor of more than 1/(1+eps). (Implying that any
//	point found subsequently in the search cannot be closer by more
//	than this factor.)
//
//	The main entry point is annkPriSearch() which sets things up and
//	then call the recursive routine ann_pri_search().  This is a
//	recursive routine which performs the processing for one node in
//	the kd-tree.  There are two versions of this virtual procedure,
//	one for splitting nodes and one for leaves. When a splitting node
//	is visited, we determine which child to continue the search on
//	(the closer one), and insert the other child into the priority
//	queue.  When a leaf is visited, we compute the distances to the
//	points in the buckets, and update information on the closest
//	points.
//
//	Some trickery is used to incrementally update the distance from
//	a kd-tree rectangle to the query point.  This comes about from
//	the fact that which each successive split, only one component
//	(along the dimension that is split) of the squared distance to
//	the child rectangle is different from the squared distance to
//	the parent rectangle.
//----------------------------------------------------------------------

//----------------------------------------------------------------------
//	To keep argument lists short, a number of global variables
//	are maintained which are common to all the recursive calls.
//	These are given below.
//----------------------------------------------------------------------

double		ANNprEps;		// the error bound
int		ANNprDim;		// dimension of space
ANNpoint	ANNprQ;			// query point
double		ANNprMaxErr;		// max tolerable squared error
ANNpointArray	ANNprPts;		// the points
ANNpr_queue	*ANNprBoxPQ;		// priority queue for boxes
ANNmin_k	*ANNprPointMK;		// set of k closest points

//----------------------------------------------------------------------
//  annkPriSearch - priority search for k nearest neighbors
//----------------------------------------------------------------------

void ANNkd_tree::annkPriSearch(
    ANNpoint		q,		// query point
    int			k,		// number of near neighbors to return
    ANNidxArray		nn_idx,		// nearest neighbor indices (returned)
    ANNdistArray	dd,		// dist to near neighbors (returned)
    double		eps)		// error bound (ignored)
{
					// max tolerable squared error
    ANNprMaxErr = ANN_POW(1.0 + eps);
    FLOP(2)				// increment floating ops

    ANNprDim = dim;			// copy arguments to static equivs
    ANNprQ = q;
    ANNprPts = pts;
    ANNptsVisited = 0;			// initialize count of points visited

    ANNprPointMK = new ANNmin_k(k);	// create set for closest k points

					// distance to root box
    ANNdist box_dist = annBoxDistance(q,
		bnd_box_lo, bnd_box_hi, dim);

    ANNprBoxPQ = new ANNpr_queue(n_pts);// create priority queue for boxes
    ANNprBoxPQ->insert(box_dist, root);	// insert root in priority queue

    while (ANNprBoxPQ->non_empty() &&
    	(!ANNmaxPtsVisited || ANNptsVisited < ANNmaxPtsVisited)) 
		{
			ANNkd_ptr np;			// next box from prior queue
			// extract closest box from queue
			ANNprBoxPQ->extr_min(box_dist, (void *&) np);
			FLOP(2)				// increment floating ops
			if (box_dist*ANNprMaxErr >= ANNprPointMK->max_key())
				break;
			np->ann_pri_search(box_dist);	// search this subtree.
    }

		for (int i = 0; i < k; i++) 
		{	// extract the k-th closest points
			dd[i] = ANNprPointMK->ith_smallest_key(i);
			nn_idx[i] = ANNprPointMK->ith_smallest_info(i);
		}

	delete ANNprPointMK;		// deallocate closest point set
	ANNprPointMK = NULL;
	delete ANNprBoxPQ;			// deallocate priority queue
	ANNprBoxPQ = NULL;
}

//----------------------------------------------------------------------
//  kd_split::ann_pri_search - search a splitting node
//----------------------------------------------------------------------

void ANNkd_split::ann_pri_search(ANNdist box_dist)
{
    ANNdist new_dist;			// distance to child visited later
					// distance to cutting plane
    ANNcoord cut_diff = ANNprQ[cut_dim] - cut_val;

    if (cut_diff < 0) {			// left of cutting plane
	ANNcoord box_diff = cd_bnds[LO] - ANNprQ[cut_dim];
	if (box_diff < 0)		// within bounds - ignore
	    box_diff = 0;
					// distance to further box
	new_dist = (ANNdist) ANN_SUM(box_dist,
		ANN_DIFF(ANN_POW(box_diff), ANN_POW(cut_diff)));

	if (child[HI] != KD_TRIVIAL)	// enqueue if not trivial
	    ANNprBoxPQ->insert(new_dist, child[HI]);
					// continue with closer child
	child[LO]->ann_pri_search(box_dist);
    }
    else {				// right of cutting plane
	ANNcoord box_diff = ANNprQ[cut_dim] - cd_bnds[HI];
	if (box_diff < 0)		// within bounds - ignore
	    box_diff = 0;
					// distance to further box
	new_dist = (ANNdist) ANN_SUM(box_dist,
		ANN_DIFF(ANN_POW(box_diff), ANN_POW(cut_diff)));

	if (child[LO] != KD_TRIVIAL)	// enqueue if not trivial
	    ANNprBoxPQ->insert(new_dist, child[LO]);
					// continue with closer child
	child[HI]->ann_pri_search(box_dist);
    }
    SPL(1)				// one more splitting node visited
    FLOP(8)				// increment floating ops
}

//----------------------------------------------------------------------
//  kd_leaf::ann_pri_search - search points in a leaf node
//
//	This is virtually identical to the ann_search for standard search.
//----------------------------------------------------------------------

void ANNkd_leaf::ann_pri_search(ANNdist box_dist)
{
    register ANNdist dist;		// distance to data point
    register ANNcoord* pp;		// data coordinate pointer
    register ANNcoord* qq;		// query coordinate pointer
    register ANNdist min_dist;		// distance to k-th closest point
    register ANNcoord t;
    register int d;

    min_dist = ANNprPointMK->max_key();	// k-th smallest distance so far

    for (int i = 0; i < n_pts; i++) 
		{	// check points in bucket

			pp = ANNprPts[bkt[i]];		// first coord of next data point
			qq = ANNprQ; 	    		// first coord of query point
			dist = 0;

			for(d = 0; d < ANNprDim; d++) 
			{
					COORD(1)			// one more coordinate hit
					FLOP(4)			// increment floating ops

					t = *(qq++) - *(pp++);	// compute length and adv coordinate
							// exceeds dist to k-th smallest?
					if( (dist = ANN_SUM(dist, ANN_POW(t))) > min_dist) 
					{
						break;
					}
			}

			if (d >= ANNprDim &&			// among the k best?
				 (ANN_ALLOW_SELF_MATCH || dist!=0)) 
			{	// and no self-match problem
								// add it to the list
					ANNprPointMK->insert(dist, bkt[i]);
					min_dist = ANNprPointMK->max_key();
			}
    }
    LEAF(1)				// one more leaf node visited
    PTS(n_pts)				// increment points visited
    ANNptsVisited += n_pts;		// increment number of points visited
}
